'''
ResNet101 图像特征提取

in: tensor(N,3,224,224)
---
out: (tensor(N,2048),tensor(N,2048,7,7))
'''

import torch
from torchvision import models
import torch.nn.functional as F


class ResNet101(torch.nn.Module):
    def __init__(self, device):
        super().__init__()
        self.checkpoint_file = 'utils/ResNet101/resnet101-63fe2227.pth'
        self.device = device
        self.model = models.resnet101()
        self.model.to(self.device)
        self.model.load_state_dict(torch.load(self.checkpoint_file, map_location=device))
        self.model.eval()

    def forward(self, x):
        x = self.model.conv1(x)
        x = self.model.bn1(x)
        x = self.model.relu(x)
        x = self.model.maxpool(x)

        x = self.model.layer1(x)
        x = self.model.layer2(x)
        x = self.model.layer3(x)
        x = self.model.layer4(x)

        feature = x.mean(3).mean(2)  # (batch_size, 2048) 是att在H和W方向上的平均
        att = F.adaptive_max_pool2d(x, (7, 7))  # (batch_size, 2048, 7, 7)

        return feature, att
